package com.gdatacloud.ai;

import ai.onnxruntime.*;
import org.opencv.core.*;
import org.opencv.imgcodecs.Imgcodecs;
import org.opencv.imgproc.Imgproc;
import org.springframework.boot.SpringApplication;
import org.springframework.boot.autoconfigure.SpringBootApplication;
import org.springframework.context.annotation.Bean;
import org.springframework.util.StringUtils;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RestController;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.net.URL;
import java.net.URLConnection;
import java.nio.FloatBuffer;
import java.util.*;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 *  环境变量(MODEL_LABEL)："fire","smoke"
 */

@RestController
@SpringBootApplication
public class DefaultAiServiceApplication {

    static  OrtEnvironment environment;

    static OrtSession session;

    static String[] names ={"angry","disgust","fear","happy","neutral","sad","surprise"};

    static int batch_size = 1;

    static int img_width = 64;

    static int img_height = 64;

    static int channel = 3;

    static int stride = 32;


    //模型输出参数数量
    long[] outputShape;

    public static void main(String[] args) {
        nu.pattern.OpenCV.loadLocally();
        SpringApplication.run(DefaultAiServiceApplication.class, args);
    }


    @Bean
    public OrtSession getOrtSession() throws OrtException, IOException {
        String MODEL_PATH = System.getenv("MODEL_PATH");
        String label = System.getenv("MODEL_LABEL");
        if(StringUtils.hasText(label)){
            names = label.split(",");
        }

        System.out.println("模型路径："+MODEL_PATH);

        environment = OrtEnvironment.getEnvironment();
        System.out.println("开始加载模型...");
        OrtSession.SessionOptions sessionOptions = new OrtSession.SessionOptions();

        if(MODEL_PATH == null){
            File onnxFile = new File("/model/model.onnx");
            if(!onnxFile.exists()){
                onnxFile = new File("/model/detect.onnx");
                if(!onnxFile.exists()){
                    throw new RuntimeException("未找到模型文件");
                }else{
                    session = environment.createSession("/model/detect.onnx", sessionOptions);
                }
            }else{
                session = environment.createSession("/model/model.onnx", sessionOptions);
            }
        }else if(MODEL_PATH.startsWith("file://")){
            String file_path = MODEL_PATH.substring(7);
            session = environment.createSession(file_path, sessionOptions);
        }else if(MODEL_PATH.startsWith("http")){
            URL url = new URL(MODEL_PATH);
            URLConnection conn = url.openConnection();
            InputStream inputStream = conn.getInputStream();
            String OS = System.getProperty("os.name").toLowerCase();

            File model_file;
            if (OS.contains("win")) {
                model_file = new File("C:/model/");
                model_file.delete();
                model_file.mkdir();
            }else{
                model_file = new File("/model/");
                model_file.delete();
                model_file.mkdir();
            }
            FileOutputStream fileOutputStream  = new FileOutputStream(model_file.getAbsoluteFile()+"/detect.onnx");
            byte[] buffer = new byte[10240];
            int i = 0;
            while (true){
                int byte_read = inputStream.read(buffer);
                if(byte_read == -1){
                    System.out.println();
                    break;
                }
                i++;
                if(i%100 == 0){
                    System.out.print("#");
                }
                fileOutputStream.write(buffer, 0, byte_read);
            }
            fileOutputStream.close();
            if (OS.contains("win")) {
                session = environment.createSession("C:/model/detect.onnx", sessionOptions);
            }else{
                session = environment.createSession("/model/detect.onnx", sessionOptions);
            }

        }else{
            if(MODEL_PATH.startsWith("s3://")||MODEL_PATH.startsWith("S3://")){
                String S3_ACCESS_KEY = System.getenv("S3_ACCESS_KEY");
                String S3_SECRET_KEY = System.getenv("S3_SECRET_KEY");
                System.out.print(S3_ACCESS_KEY+" : "+S3_SECRET_KEY);
                session = environment.createSession(MODEL_PATH, sessionOptions);
            }else{
                throw new RuntimeException("MODEL_PATH模型路径环境变量值错误");
            }
        }
        System.out.println("模型加载成功，开始初始化...");
        if(names == null || names.length == 0){
            String class_name = session.getMetadata().getCustomMetadataValue("classes").orElse(null);
            if(!StringUtils.hasText(class_name)){
                class_name = session.getMetadata().getCustomMetadataValue("names").orElse(null);
                if(!StringUtils.hasText(class_name)){
                    if(System.getenv("ClASS_NAME") == null){
                        throw new RuntimeException("缺少模型标签");
                    }else{
                        class_name = System.getenv("ClASS_NAME");
                    }
                }
            }
            names = new String[class_name.split(",").length];
            Pattern pattern = Pattern.compile("'([^']*)'");
            Matcher matcher = pattern.matcher(class_name);

            int k = 0;
            while (matcher.find()) {
                names[k] = matcher.group(1);
                k++;
            }
        }

        System.out.println("标签数量： " + names.length);
        System.out.println("标签： " + Arrays.toString(names));
        System.out.println("模型输入tensor：");
        session.getInputInfo().keySet().forEach(key -> {
            try {
                TensorInfo input = (TensorInfo) session.getInputInfo().get(key).getInfo();
                System.out.println(key+" : "+ Arrays.toString(input.getShape()));
            } catch (OrtException e) {
                throw new RuntimeException(e);
            }
        });

        int i = 0;
        System.out.println("模型输出tensor：");
        for (String key : session.getOutputInfo().keySet()){
            try {
                TensorInfo output = (TensorInfo) session.getOutputInfo().get(key).getInfo();
                if(i == 0){
                    outputShape = output.getShape();
                }
                System.out.println(key+" : "+ Arrays.toString(output.getShape()));
                i++;
            } catch (OrtException e) {
                throw new RuntimeException(e);
            }
        }

        System.out.println("模型初始化成功，启动web服务...");
        System.out.println("接口地址 POST ：http://127.0.0.1:8080/detect");
        System.out.println("请求类型： Content-Type: application/json");
        System.out.println("参数： { image : 图像base64字符串 }");
        return session ;
    }


    // 去掉base64逗号之前的头部信息
    @PostMapping("detect")
    public Object detect(@RequestBody Input input) {
        if(input == null || input.getImage() == null){
            return  ImgUtils.response(3, "缺少image参数", null);
        }

        int index = input.getImage().indexOf(",");
        if(index != -1){
            input.setImage(input.getImage().substring(index+1));
        }

        try {
            Mat img = Imgcodecs.imdecode(new MatOfByte(Base64.getDecoder().decode(input.getImage())), Imgcodecs.IMREAD_COLOR);
            Letterbox letterbox = new Letterbox();
            img = letterbox.letterbox(img);
            Imgproc.cvtColor(img, img, Imgproc.COLOR_BGR2RGB);
            img.convertTo(img, CvType.CV_32FC1, 1. / 255);
            float[] whc = new float[batch_size * channel * img_width * img_height];
            img.get(0, 0, whc);
            float[] chw =ImgUtils. whc2cwh(whc);
            FloatBuffer inputBuffer = FloatBuffer.wrap(chw);
            OnnxTensor tensor = OnnxTensor.createTensor(environment, inputBuffer, new long[]{batch_size , channel , img_width , img_height});
            HashMap<String, OnnxTensor> stringOnnxTensorHashMap = new HashMap<>();
            stringOnnxTensorHashMap.put(session.getInputInfo().keySet().iterator().next(), tensor);
            OrtSession.Result output = session.run(stringOnnxTensorHashMap);

            TensorInfo resultTensorInfo = (TensorInfo) output.get(0).getInfo();
            List<ClassificationOutput> results = new ArrayList<>();


            if(resultTensorInfo.getShape()[0] == 0){
                return  ImgUtils.response(0, "成功！", new ArrayList<>());
            }
            float[][] outputData = (float[][]) output.get(0).getValue();
            for(float[] x : outputData){
                for (int i = 0; i < x.length; i++) {
                    results.add(new ClassificationOutput(names[i],x[i]));
                }
            }
            System.out.println(new Date() + " : "+results);
            results.sort(Comparator.comparing(ClassificationOutput::getScore).reversed());
            return  ImgUtils.response(0, "成功！", results);
        } catch (Exception e){
            System.out.println(e.getMessage());
            return ImgUtils.response(3, "失败！参数异常", null);
        }
    }



}

